import os
from pathlib import Path

from flask import Flask, jsonify, request
from exception.T2Exception import T2Exception
from util.T2es import T2es
from util.T2llm import T2llm

"""
知识库相关接口
"""
app = Flask(__name__)

#知识库index
RAG_KNOWLEDGE_INDEX = "rag_knowledge_index"


def json(data=None,  msg="OK", code=1):
    return jsonify({
        "code": code,
        "msg": msg,
        "data": data
    }), 200

@app.errorhandler(Exception)
def handle_exception(e):
    print(e)
    return json(str(e), "系统错误", -1)

@app.errorhandler(T2Exception)
def handle_exception(e):
    return json(str(e), "业务异常", 0)


@app.route('/')
def index():
    return json("欢迎来到我的Flask应用！")


def create_index():
    # 创建索引
    mappings = {
        "mappings" : {
          "properties" : {
              "kb_id": {  #知识库ID，方便直接针对某一类文档进行查询
                  "type": "keyword"
              },
              "file_name": {
                  "type": "keyword"
              },
              "file_id": {    #文档id 方便直接针对某一个文档查询
                  "type": "keyword"
              },
              "content": {
                  "type": "text"
              },
              "content_vector" : {
                  "type" : "dense_vector",
                  "dims" : 1024  #需要embedding_model比对一致
              }
          }
        }
    }
    es = T2es()
    es.create_index(RAG_KNOWLEDGE_INDEX, mappings);
    es.close()

# 文档上传、切分、向量化、保存
@app.route('/upload', methods=['POST'])
def upload():
    # 检查是否有文件上传
    if 'file' not in request.files:
        raise T2Exception('没有选择文件')
    file = request.files['file']
    # 如果用户没有选择文件，浏览器也会提交一个没有文件名的空部分
    if file.filename == '':
        raise T2Exception('没有选择文件')

    # 确保保存文件的目录存在
    upload_dir = os.path.join(Path(__file__).parent,'uploads')
    if not os.path.exists(upload_dir):
        os.makedirs(upload_dir)

    # 构建文件保存路径
    filepath = os.path.join(upload_dir, file.filename)
    file.save(filepath)

    doc_info = T2llm().embed_file(filepath)

    kb_id = request.form.get("kb_id", "common")

    # es入库
    es = T2es()
    es.add_vector_docs(RAG_KNOWLEDGE_INDEX, kb_id, file.filename, doc_info.get("texts"), doc_info.get("vectors"))
    es.close()

    return json()

"""
知识库查询
query:问答文本
kb_ids?: [] 知识库ids,如果不传则查询所有
score_min_threshold?: 最小相似度
size_max_threshold?：阈值，最大记录行数
"""
@app.route('/query', methods=['POST'])
def query():
    data = request.get_json()
    text =  str(data.get('text', ''))
    kb_ids = list(data.get('kb_ids', []))
    file_id = data.get('file_id')
    score_min_threshold = float(data.get('score_min_threshold', 0.5))
    size_max_threshold = int(data.get('size_max_threshold', 10))


    t2llm = T2llm()
    vector = t2llm.embed_text(text)

    es = T2es()
    results = es.get_vector_docs(RAG_KNOWLEDGE_INDEX,vector,kb_ids,score_min_threshold, size_max_threshold, file_id)
    es.close()




    return json(results)


if __name__ == '__main__':
    create_index()
    app.run(port=5001)
